Add limit to ArrowReaderBuilder to push limit down to parquet reader #3633
Conversation
| Self { | ||
| selectors: vec![ | ||
| RowSelector::select(limit), | ||
| RowSelector::skip(total_rows.saturating_sub(limit)), |
There was a problem hiding this comment.
| RowSelector::skip(total_rows.saturating_sub(limit)), |
This shouldn't be necessary, it will get removed by RowSelection::trim anyway. Does make me wonder if this method is even really needed 🤔
There was a problem hiding this comment.
Cool, I think the main thing it does is just to truncate a final select if needed, but it can be a lot simpler if we do not need to append a skip on the end
There was a problem hiding this comment.
Oh, sorry wrong method, yeah I don't think it makes sense then
| @@ -371,6 +382,35 @@ impl RowSelection { | |||
| self | |||
| } | |||
|
|
|||
| /// Limit this [`RowSelection`] to only select `limit` rows | |||
| pub(crate) fn limit(mut self, mut limit: usize) -> Self { | |||
| let mut remaining = 0; | |||
There was a problem hiding this comment.
This could be more simply implemented as
Self {
selectors: intersect_row_selections(&self.selectors, &[RowSelector::select(limit)]),
}
There was a problem hiding this comment.
Yes, RowSelection row count allow to less than file total row count 🤔.
There was a problem hiding this comment.
intersect_row_selections is a bit weird. It doesn't seem to coalesce selections. Not sure if that is intentional or not.
| @@ -453,6 +467,17 @@ impl<T: ChunkReader + 'static> ArrowReaderBuilder<SyncReader<T>> { | |||
| selection = Some(RowSelection::from(vec![])); | |||
| } | |||
|
|
|||
| // If a limit is defined, apply it to the final `RowSelection` | |||
There was a problem hiding this comment.
We need to also do something similar for the async reader
Co-authored-by: Raphael Taylor-Davies <1781103+tustvold@users.noreply.github.com>
| @@ -430,6 +432,17 @@ where | |||
| return Ok((self, None)); | |||
| } | |||
|
|
|||
| // If a limit is defined, apply it to the final `RowSelection` | |||
| if let Some(limit) = limit { | |||
There was a problem hiding this comment.
I think this will apply the limit to each row group separately, instead of the entire file?
There was a problem hiding this comment.
Yeah, you're right. Fixed this to track the limit across row groups.
|
Thank you |
|
Benchmark runs are scheduled for baseline = 3057fa5 and contender = a76ea1c. a76ea1c is a master commit associated with this PR. Results will be available as each benchmark for each run completes. |
Which issue does this PR close?
Closes #3631
Rationale for this change
What changes are included in this PR?
Are there any user-facing changes?